import torch
from torch.utils.data import Dataset, TensorDataset, random_split, DataLoader
import numpy as np


class TitanicDataset(Dataset):
    def __init__(self, filepath):
        XY_data = np.loadtxt(filepath, delimiter=',', dtype=np.float32, skiprows=1)
        self.len = XY_data.shape[0]
        self.X = torch.from_numpy(XY_data[:, 2:])
        self.y = torch.from_numpy(XY_data[:, [1]])

    def __getitem__(self, index):
        return self.X[index], self.y[index]

    def __len__(self):
        return self.len


titanicDataset = TitanicDataset("../data/Titanic/titanic_preprocessed3.csv")
print(titanicDataset.__getitem__(0))

dataset = TensorDataset(titanicDataset.X, titanicDataset.y)
print(dataset.tensors[0])

train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, pin_memory=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, pin_memory=True)
for i, (inputs, labels) in enumerate(train_loader):
    print(inputs.shape[0])
    print(labels.shape)
    break

pred = torch.tensor([0.8, 0.2, 0.6])  # 预测概率
labels = torch.tensor([1.0, 0.0, 0.0])  # 标签

mask = pred > 0.5
print("pred > 0.5 =", mask)  # tensor([True, False, True])

comp = (mask == labels)
print("比较结果 =", comp)  # tensor([ True,  True, False])
